package airthmetic.exercise.trie;

// 实现前缀树
public class Trie {
    class TrieNode {
        private boolean isEnd;
        private int R = 26;
        private TrieNode[] tns;
        private String word;

        public TrieNode() {
            tns = new TrieNode[R];
        }

        public boolean contains(char ch) {
            return tns[ch - 'a'] != null;
        }

        public void put(char ch, TrieNode node) {
            tns[ch - 'a'] = node;
        }

        public TrieNode get(char ch) {
            return tns[ch - 'a'];
        }

        public void setEnd() {
            isEnd = true;
        }

        public boolean isEnd() {
            return isEnd;
        }

    }

    TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String word) {
        if (word == null || word.length() == 0) {
            return;
        }
        char[] wordArr = word.toCharArray();
        TrieNode node = root;
        for (char w : wordArr) {
            if (!node.contains(w)) {
                node.put(w, new TrieNode());
            }
            node = node.get(w);
        }
        node.setEnd();
        node.word = word;
    }

    public boolean search(String word) {
        if (word == null || word.length() == 0) {
            return false;
        }
        TrieNode node = searchPrefix(word);
        return node != null && node.isEnd();
    }

    public TrieNode searchPrefix(String word) {
        if (word == null || word.length() == 0) {
            return null;
        }
        char[] wordArr = word.toCharArray();
        TrieNode node = root;
        for (char w : wordArr) {
            if (!node.contains(w)) {
                return null;
            }
            node = node.get(w);
        }

        return node;
    }

    public boolean startsWith(String prefix) {
        if (prefix == null || prefix.length() == 0) {
            return false;
        }
        TrieNode node = searchPrefix(prefix);
        return node != null;
    }
}
